MAYAVI_EXIST = True
try:
    from mayavi import mlab
except ModuleNotFoundError as e:
    if e.name == 'mayavi':
        MAYAVI_EXIST = False
    else:
        raise e
import numpy as np
from numpy.typing import NDArray

from hj_dataset_devkit.common_util import transform_3d_cloud

from common_utils.math_util import euler_to_rotation

class Colors:
    Red = (0.9, 0.3, 0.24)
    Orange = (0.9, 0.49, 0.13)
    Yellow = (0.95, 0.8, 0.25)
    Green = (0.32, 0.8, 0.5)
    Cyan = (0.3, 0.82, 0.88)
    Blue = (0.3, 0.63, 0.9)
    Purple = (0.65, 0.41, 0.75)
    White = (1., 1., 1.)
    Black = (0., 0., 0.)
    Dark = (0.25, 0.25, 0.25)

    @staticmethod
    def random(seed=None):
        import random
        random.seed(seed)
        return (random.random(), random.random(), random.random())

def show_text(figure, text: str, point: np.ndarray, rgb=Colors.White, scale=1.):
    assert point.size > 2
    mlab.text3d(point[0], point[1], point[2], text, color=rgb, scale=scale, figure=figure)

def show_line(figure, points: NDArray, rgb=Colors.White, scale=1.):
    assert points.shape[1] > 2
    mlab.plot3d(points[:, 0], points[:, 1], points[:, 2],
                color=rgb, tube_radius=None, line_width=scale, figure=figure)

def show_box(figure, center, size, rotation, rgb=Colors.White, scale=1.) -> None:
    assert len(center) == len(size) == len(rotation) == 3
    box_pos, box_size, box_rot = np.array(center), np.array(size), np.array(rotation)
    pts_prefix = np.array([[-1, -1, -1], [-1, 1, -1], [1, 1, -1], [1, -1, -1],
                           [-1, -1, 1], [-1, 1, 1], [1, 1, 1], [1, -1, 1]]) / 2
    box_pts = euler_to_rotation(box_rot).apply(pts_prefix * box_size) + box_pos
    # plot bottom and upper frame, then 4 columns
    plot_order = [0, 1, 2, 3, 0, 4, 5, 6, 7, 4]
    mlab.plot3d(box_pts[:, 0][plot_order], box_pts[:, 1][plot_order], box_pts[:, 2][plot_order],
                color=rgb, tube_radius=None, line_width=scale, figure=figure)
    for i in range(1, 4):
        plot_order = [i, i + 4]
        mlab.plot3d(box_pts[:, 0][plot_order], box_pts[:, 1][plot_order], box_pts[:, 2][plot_order],
                    color=rgb, tube_radius=None, line_width=scale, figure=figure)

def show_coordinate(figure, pose: NDArray=np.identity(4), length=1., scale=3.) -> None:
    assert pose.shape == (4, 4)
    assert length > 0
    coors = np.array([[0, 0, 0], [length, 0, 0], [0, length, 0], [0, 0, length]], dtype=np.float64)
    coors = transform_3d_cloud(coors, pose[:3, 3], pose[:3, :3])
    show_line(figure, coors[[0, 1]], Colors.Red, scale)
    show_line(figure, coors[[0, 2]], Colors.Green, scale)
    show_line(figure, coors[[0, 3]], Colors.Blue, scale)

def show_np_cloud(figure, cloud: np.ndarray, rgb=Colors.White, scale=0.05, mode='point'):
        assert cloud.shape[1] > 2
        if len(rgb) == 3 and min(rgb) >= 0 and max(rgb) <= 1:
            mlab.points3d(cloud[:, 0], cloud[:, 1], cloud[:, 2], mode=mode, color=rgb, scale_factor=scale, figure=figure)
        else:
            mlab.points3d(cloud[:, 0], cloud[:, 1], cloud[:, 2], rgb, mode=mode, opacity=0.6,
                          colormap='jet', scale_factor=scale, scale_mode='none', figure=figure)
